package cn.fzkj.ailangchain4j.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

/**
 * @ description
 * @ author yaya
 * @ since 2025/11/4
 */
@Configuration
public class RagConfig {

    @Autowired
    private EmbeddingModel embeddingModel;
//    @Autowired
    private EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

    @Bean
    public ContentRetriever contentRetriever() {
        // 加载文档
        List<Document> document = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs");
        // 切割文档,最多120字符重叠
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(1000, 200);
        // 定义文档加载器，把文档转换为向量存储到向量数据库中
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(splitter)
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                // 为了提高文档质量，将切割后的文档碎片 TextSegment 拼接上文档名称作为元信息
                .textSegmentTransformer(segment -> {
                    return TextSegment.from(segment.metadata().getString("file_name")
                            + "\n" + segment.text(), segment.metadata());
                })
                .build();
        ingestor.ingest(document);
        return EmbeddingStoreContentRetriever.builder()
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .maxResults(5)
                .minScore(0.75) // 置信度
                .build();
    }
}
